import json

import orjson

from time import time
from pathlib import Path
from random import sample
from math import cos, sin, radians
from typing import Dict, List, Union, Tuple, Callable, Any
from functools import wraps

from .decorator import param2path
from .img import ImageTools
from .label_me import LabelMe
from .log import logger
from .utils import P_TYPE


def verify_json(f_d: Dict, l_d: Dict = None) -> Callable:
    """
    校验文件是否被标注
    Args:
        f_d: 图片参数信息
        l_d: 标注文件参数信息

    Returns:

    """

    def wrapper(func: Callable) -> Callable:
        @wraps(func)
        def decorator(*args, **kwargs) -> Any:

            def _verify(f: Path, f_l: Path = None) -> bool:
                if _f_l is None:
                    _label_path = Path(f.parent)
                else:
                    _label_path = f_l
                _ret = _label_path.joinpath(f"{f.stem}.json").exists()
                if not _ret:
                    logger.warning(f"{f.name} 没有对应的标注文件")

                return _ret

            _f = _f_l = None
            for _k, _i in f_d.items():
                if _i is not None and _i < len(args):
                    _f = args[_i]
                if _k in kwargs:
                    _f = kwargs[_k]

            for _k_l, _i_l in l_d.items():
                if _i_l is not None and _i_l < len(args):
                    _f_l = args[_i_l]
                if _k_l in kwargs:
                    _f_l = kwargs[_k_l]
            result = _verify(_f, _f_l)

            if result:
                return func(*args, **kwargs)
            else:
                return

        return decorator

    return wrapper


@param2path({'f': 0})
def rotate(f: Union[Path, str], angle: int, is_show: bool = False,
           to_path: str = r'../rotate') -> None:
    """
    label me 多边形标注图片数据旋转
    Args:
        f: 已标注图片地址
        angle: 旋转角度
        is_show: 是否展示图片标注后样式
        to_path: 新数据存放目录，相对于文件目录

    Returns:

    """
    base_path_c = Path(f.parent)  # 文件所在目录
    to_path_c = base_path_c.joinpath(to_path)  # 输出路径
    if not to_path_c.exists():
        to_path_c.mkdir(parents=True, exist_ok=True)
    # 判断标注文件是否存在
    _file_c = base_path_c.joinpath(f"{f.stem}.json")
    if _file_c.exists():
        _stem = f"{f.stem}_{int(time() * 1000)}"

        _img_tools = ImageTools(f)

        _img_tools.rotate(angle)  # 图片旋转

        # 获取标注文件内容
        _json = orjson.loads(_file_c.read_text(encoding='utf-8'))  # type: Dict
        _shapes = _json.get('shapes')  # type: List
        _x0, _y0 = _img_tools.center_rotate  # 旋转中心
        shapes = []  # 形状
        for _shape in _shapes:  # type: Dict
            _points = _shape.get('points')  # type: List
            points = []  # 坐标点
            # 计算旋转后坐标
            for _x, _y in _points:
                # 以旋转点为原点偏移坐标
                tx = (_x - _x0)
                ty = (_y - _y0)
                _alpha = radians(angle)  # 弧度制
                _offset = _img_tools.center_offset  # 旋转后图片中心偏移量
                x1 = cos(_alpha) * tx + sin(_alpha) * ty + _x0 + _offset[0]
                y1 = -sin(_alpha) * tx + cos(_alpha) * ty + _y0 + _offset[1]
                points.append((x1, y1))  # 记录
            _shape['points'] = points  # 回填标注点坐标
            shapes.append(points)

        # 回填其他信息
        _json['imagePath'] = f"{_stem}{f.suffix}"
        _json['imageData'] = _img_tools.to_base64
        _json['imageWidth'], _json['imageHeight'] = _img_tools.size
        # 保存旋转图片 没有文件夹时自动生成
        _img_tools.save(str(to_path_c), _stem)
        # 保存标注文件
        to_file_c = to_path_c.joinpath(f"{_stem}.json")
        logger.info(f"旋转后标注文件地址：{to_file_c}")
        to_file_c.write_text(
            json.dumps(_json, ensure_ascii=False, indent=2), encoding='utf-8')
        # 绘制坐标并展示
        if is_show:
            (_img_tools
             .draw_point([j for i in shapes for j in i])
             .show('plt')
             )
    else:
        logger.warning(f"没有对应的标注文件")


@param2path({'path_c': 0})
def random_rotate(path_c: Union[Path, str], k: int, angle: int = None,
                  angle_zoom: float = .5, angle_start: int = None,
                  **kwargs) -> None:
    """
    随机旋转图片
    Args:
        k: 增广数
        path_c: 图片文件夹
        angle: 角度步长
        angle_zoom: 角度步长 缩放比例（不能大于1）
        angle_start: 开始角度
        **kwargs: rotate() 参数

    Returns:

    """
    if not angle:
        angle = int(360 // k * angle_zoom)
    if angle_start is None or angle_start < 0:
        angle_start = angle
    k_max_ = (360 - angle_start) // angle - 1  # 计算最大 k_ 值
    _start_max = 360 - angle * 2 + 1  # 最大开始角度

    assert k_max_ > 0, (f"k_max_ 必须大于 0"
                        f"，请修改 angle_start angle 参数"
                        f"，保证 (360 - angle_start) // angle - 1 > 0")
    assert k_max_ >= k, f"k_ 值不能大于 {k_max_}"
    assert 0 < angle_zoom <= 1, f"angle_scale 不能大于 1 或 小等于 0"
    assert angle >= 5, f"angle 的值不能小于 5 ，当前为：{angle}"
    assert angle_start < _start_max, f"angle_start 必须小于 {_start_max}"

    # 数据旋转增广
    for f in path_c.glob('*.jpg'):
        logger.info(f"待旋转文件地址：{path_c}")
        for _, _angle in enumerate(
                sample(range(angle_start, 360 - angle + 1, angle), k=k)
        ):
            logger.info(f"随机旋转角度：{_angle}")
            rotate(f, _angle, **kwargs)


@param2path({'file_path_c': 0})
def polygon2rectangle(file_path_c: Union[Path, str],
                      to_path: str = r'../rectangle') -> None:
    """
    多边形标注信息转矩形标注信息
    Args:
        file_path_c: 标注文件地址
        to_path: 新数据存放目录，相对于文件目录

    Returns:

    """
    # 标注文件数据
    _json = orjson.loads(file_path_c.read_text(encoding='utf-8'))

    to_path_c = Path(file_path_c.parent).joinpath(to_path)
    if not to_path_c.exists():
        to_path_c.mkdir(parents=True, exist_ok=True)

    # 遍历标注坐标数据并修改
    for _shape in _json.get('shapes'):
        points = []  # 新坐标
        if (_t := _shape['shape_type']) != 'polygon':
            logger.warning(f"标注框类型不正确，期望：rectangle，实际：{_t}")
            continue
        # 获取标注点数据
        _xs, _ys = list(zip(*_shape['points']))  # type: List
        points.append((min(_xs), min(_ys)))
        points.append((max(_xs), max(_ys)))
        # 数据回填
        _shape['points'] = points
        _shape['shape_type'] = 'rectangle'

    # 保存标注文件
    to_file_c = to_path_c.joinpath(f"{file_path_c.name}")
    logger.info(f"矩形标注文件地址：{to_file_c}")
    to_file_c.write_text(
        json.dumps(_json, ensure_ascii=False, indent=2), encoding='utf-8')


@param2path({'base_path_c': 0})
def polygon2rectangle_dir(base_path_c: Union[Path, str], **kwargs) -> None:
    """
    转换文件夹下所有多变形标注数据
    Args:
        base_path_c: 标注文件夹地址
        **kwargs: polygon2rectangle() 参数

    Returns:

    """
    for f in base_path_c.glob('*.json'):
        logger.info(f"待转换多边形标注文件地址：{f}")
        polygon2rectangle(f, **kwargs)


@param2path({'f': 0, 'to_path_c': -2, 'label_path': -1}, ['to_path_c'])
@verify_json({'f': 0}, {'label_path': -1})
def resize(f: Union[Path, str], ratio: float = 1, max_side: int = 0,
           min_side: int = 0, d_size: Tuple = None,
           to_path_c: Union[Path, str] = r'zoom',
           label_path: P_TYPE = None,
           ) -> None:
    """
    图片数据尺寸修改
    Args:
        f: 已标注文件地址
        ratio: 比例
        max_side: 最大边长
        min_side: 最小边长
        d_size: 指定尺寸
        to_path_c: 存放路径

    Returns:

    """
    # 修改图片尺寸
    _img_tools = ImageTools(f)
    _w_o, _h_o = _img_tools.size

    (_img_tools
     .resize(d_size=d_size, max_side=max_side, min_side=min_side, ratio=ratio)
     .save(to_path_c.joinpath('images')))

    # 获取新的尺寸
    _w, _h = _img_tools.size

    if label_path is None:
        label_path = Path(f.parent)
    _label_me = LabelMe.from_json(label_path.joinpath(f"{f.stem}.json"))

    # 遍历标注坐标数据并修改
    for _shape in _label_me.shapes:

        for _point in _shape.points:
            _point.x = _point.x * _w / _w_o
            _point.y = _point.y * _h / _h_o

    # 回填其他信息
    _label_me.image_base64 = _img_tools.to_base64
    _label_me.width, _label_me.height = _w, _h

    # 保存标注文件
    logger.info(f"新尺寸标注文件地址：{to_path_c}")
    _label_me.save(to_path_c.joinpath('annotations'))


@param2path({'base_path_c': 0})
def resize_dir(base_path_c: Union[Path, str], **kwargs) -> None:
    """
    转换文件夹下所有标注数据的尺寸
    Args:
        base_path_c: 标注文件夹地址
        **kwargs: resize() 参数

    Returns:

    """
    for f in base_path_c.glob('*.jpg'):
        logger.info(f"修改尺寸地址：{f}")
        resize(f, **kwargs)


@param2path({'f': 0, 'to_path_c': 1, 'label_path': 3}, ['to_path_c'])
@verify_json({'f': 0}, {'label_path': 3})
def fill(f: P_TYPE, to_path_c: P_TYPE, size: Union[Tuple, float],
         label_path: P_TYPE = None, **kwargs) -> None:
    """
    图片填充
    Args:
        f: 图片路径
        to_path_c: 输出路径
        kwargs: 填充方法参数

    Returns:

    """
    if label_path is None:
        label_path = Path(f.parent)
    label_file_path = label_path.joinpath(f"{f.stem}.json")
    # 填充及保存 图片
    _image_tools = ImageTools(f)
    (_image_tools.fill(size).save(to_path_c.joinpath('images')))
    _t, _b, _l, _r = _image_tools.fill_num
    # 处理 标注文件
    _label_me = LabelMe.from_json(label_file_path)
    _label_me.image_base64 = _image_tools.to_base64
    _label_me.width, _label_me.height = _image_tools.size
    _shapes = []
    for _j, _shape in enumerate(_label_me.shapes):
        for _i, _point in enumerate(_shape.points):
            _point.x += _l
            _point.y += _t
        _label_me.shapes[_j] = _shape
    # 保存 标注文件
    _label_me.save(to_path_c.joinpath('annotations'))
    # _to_path_c_label = to_path_c.joinpath('annotations')
    # if not _to_path_c_label.exists():
    #     _to_path_c_label.mkdir(parents=True, exist_ok=True)
    # _to_path_c_label.joinpath(f"{f.stem}.json").write_text(
    #     json.dumps(_label_me.to_dict, ensure_ascii=False, indent=2),
    #     encoding='utf-8'
    # )
